
"""
获取数据表
"""
from core.data_analyze.state import AnalyzeState
from langgraph.types import Command
from langchain_core.messages import AIMessage, ToolMessage
from langgraph.config import get_stream_writer
from core.data_analyze.utils import service_provider
from loguru import logger


def fetch_data_node(state: AnalyzeState):

    writer = get_stream_writer()
    writer("## **获取相应的数据表** ")
    logger.info("## **获取相应的数据表** ")

    query = state.get('task_description')
    revised_task = state.get('revised_task')
    if revised_task is not None and revised_task != '':
        query = revised_task
    try:
        datasets = service_provider.get_data(query)
    except Exception as err:
        return Command(
            update={
                "messages": [ToolMessage(content=str(err), tool_name='service_provider.get_data')],
            },
            goto="human_feedback"
        )
    if isinstance(datasets, str):
        return Command(
            update={
                "messages": [AIMessage(content=datasets, name="fetch_data")]
            },
            goto="human_feedback"
        )
    elif isinstance(datasets, list) and len(datasets) != 0:
        return Command(
            update={
                'revised_task': query,
                "datasets": datasets,
                "cur_retry_times": 0
            },
            goto="task_decompose",
        )
    else:
        writer("未获取到有效数据.")
        logger.info("未获取到有效数据.")
        max_retry_times = state.get('max_retry_times')
        cur_retry_times = state.get('cur_retry_times')
        if cur_retry_times < max_retry_times:
            return Command(
                update={
                    "cur_retry_times": cur_retry_times + 1
                },
                goto="task_rewrite",
            )
        else:
            writer("达到最大重试次数.")
            logger.info("达到最大重试次数.")
            return Command(goto="human_feedback")
